from absl import app, flags
import collections
import h5py
import os
import pandas
import socket
import traceback
import re

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import coin.strategy.mm.tool.archive_base as abase
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2


_fi_repo = '/remote/iosg/data-2/buckets/feed.derived.interval_h5'
_fi_root = f'{_fi_repo}/coin/main'
_fi_resol = 'PT1M'
fi_dir = f'{_fi_root}/{_fi_resol}'


def load_mea(mea, trading_date, symbol=None):
  tdstr = trading_date.strftime("%Y%m%d")
  dfdict = {}
  for h5path in [
      f'{fi_dir}/{mea}/{tdstr}/{mea}--ohlc.h5',
      f'{fi_dir}/{mea}/{tdstr}/{mea}--volume.h5']:
    try:
      h5data = h5py.File(h5path, 'r')
    except:
      traceback.print_exc()
      raise ValueError(h5path)
    nativesymbolstrs = h5data['universe'][:]
    mt, ex, api = mea.split(".")
    cols = []
    chosen = []
    for i, nativesymbolstr in enumerate(nativesymbolstrs):
      if symbol is not None and not re.match(symbol, nativesymbolstr):
        continue
      chosen.append(i)
      try:
        product = generate_product_from_str2(mt, ex, api, nativesymbolstr, trading_date)
        cols.append(
            product.subscription_symbol.replace("CURRENT_", "").replace("THIS_", "")
            if hasattr(product, 'subscription_symbol')
            else product.symbol)
      except Exception:
        cols.append(nativesymbolstr)
    tidx = pandas.DatetimeIndex(h5data['timestamp'][:])
    for key in h5data.keys():
      if key not in ['timestamp', 'universe']:
        df = pandas.DataFrame(h5data[key][:,chosen])
        df.index = tidx
        df.columns = cols
        dfdict[key] = df
  return dfdict


def load_df_dict(mea, trading_date, symbol=None):
  trading_dates = abase.get_trading_dates(trading_date)
  dfs_dict = collections.defaultdict(list)
  for trading_date in trading_dates:
    df_dict_each = load_mea(mea, trading_date, symbol=symbol)
    for df_key, df in df_dict_each.items():
      dfs_dict[df_key].append(df)

  df_dict = {}
  for key, dfs in dfs_dict.items():
    df_dict[key] = pandas.concat(dfs, axis=0, sort=False)
  return df_dict

def main(_):
  df_dict = load_df_dict(flags.FLAGS.mea, flags.FLAGS.trading_date)
  if flags.FLAGS.mode == 'pdb':
    import pdb
    pdb.set_trace()


if __name__ == "__main__":
  flags.DEFINE_string('mea', 'Spot.Upbit.v1', '')
  flags.DEFINE_string('mode', 'pdb', 'plot,pdb')
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  app.run(main)
